"""
回测结果可视化模块

提供回测结果的可视化功能，包括：
- 资金曲线图
- 回撤曲线图
- 交易信号标注
- 绩效指标展示

作者: AI Assistant
版本: 1.0.0
日期: 2025-01-06
"""

import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import streamlit as st
from typing import Dict, List, Optional


class BacktestVisualizer:
    """
    回测可视化器
    
    提供多种图表来展示回测结果。
    
    属性:
        result: 回测结果字典
        data: 原始K线数据
        strategy: 策略实例
    """
    
    def __init__(self, result: Dict, data: pd.DataFrame, strategy=None):
        """
        初始化可视化器
        
        参数:
            result: 回测结果字典
            data: 原始K线数据
            strategy: 策略实例（可选，用于获取策略指标）
        """
        self.result = result
        self.data = data
        self.strategy = strategy
    
    def plot_equity_curve(self) -> go.Figure:
        """
        绘制资金曲线图
        
        返回:
            go.Figure: Plotly图表对象
        """
        equity_df = self.result['equity_curve']
        
        fig = go.Figure()
        
        # 添加总资产曲线
        fig.add_trace(go.Scatter(
            x=equity_df.index,
            y=equity_df['总资产'],
            name='总资产',
            line=dict(color='#2E86DE', width=2),
            hovertemplate='日期: %{x}<br>总资产: ¥%{y:,.2f}<extra></extra>'
        ))
        
        # 添加现金曲线
        fig.add_trace(go.Scatter(
            x=equity_df.index,
            y=equity_df['现金'],
            name='现金',
            line=dict(color='#54A0FF', width=1, dash='dot'),
            hovertemplate='日期: %{x}<br>现金: ¥%{y:,.2f}<extra></extra>'
        ))
        
        # 添加持仓市值曲线
        fig.add_trace(go.Scatter(
            x=equity_df.index,
            y=equity_df['持仓市值'],
            name='持仓市值',
            line=dict(color='#FF9FF3', width=1, dash='dot'),
            fill='tonexty',
            fillcolor='rgba(255, 159, 243, 0.1)',
            hovertemplate='日期: %{x}<br>持仓市值: ¥%{y:,.2f}<extra></extra>'
        ))
        
        # 更新布局
        fig.update_layout(
            title='资金曲线',
            xaxis_title='日期',
            yaxis_title='金额（元）',
            hovermode='x unified',
            template='plotly_white',
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        
        return fig
    
    def plot_drawdown(self) -> go.Figure:
        """
        绘制回撤曲线图
        
        返回:
            go.Figure: Plotly图表对象
        """
        equity_df = self.result['equity_curve']
        equity_series = equity_df['总资产']
        
        # 计算回撤
        running_max = equity_series.expanding().max()
        drawdown = (equity_series - running_max) / running_max * 100
        
        fig = go.Figure()
        
        # 添加回撤曲线
        fig.add_trace(go.Scatter(
            x=drawdown.index,
            y=drawdown,
            name='回撤',
            fill='tozeroy',
            fillcolor='rgba(239, 83, 80, 0.3)',
            line=dict(color='#EF5350', width=2),
            hovertemplate='日期: %{x}<br>回撤: %{y:.2f}%<extra></extra>'
        ))
        
        # 标注最大回撤点
        max_dd_date = drawdown.idxmin()
        max_dd_value = drawdown.min()
        
        fig.add_annotation(
            x=max_dd_date,
            y=max_dd_value,
            text=f'最大回撤: {max_dd_value:.2f}%',
            showarrow=True,
            arrowhead=2,
            arrowcolor='red',
            font=dict(color='red', size=12)
        )
        
        # 更新布局
        fig.update_layout(
            title='回撤曲线',
            xaxis_title='日期',
            yaxis_title='回撤 (%)',
            hovermode='x unified',
            template='plotly_white',
            height=400
        )
        
        return fig
    
    def plot_price_with_signals(self) -> go.Figure:
        """
        绘制价格图并标注交易信号
        
        返回:
            go.Figure: Plotly图表对象
        """
        # 创建子图
        fig = make_subplots(
            rows=2, cols=1,
            shared_xaxes=True,
            vertical_spacing=0.03,
            subplot_titles=('价格与交易信号', '持仓数量'),
            row_heights=[0.7, 0.3]
        )
        
        # 1. K线图
        fig.add_trace(go.Candlestick(
            x=self.data.index,
            open=self.data['开盘'],
            high=self.data['最高'],
            low=self.data['最低'],
            close=self.data['收盘'],
            name='K线',
            increasing_line_color='#EF5350',
            decreasing_line_color='#26A69A',
        ), row=1, col=1)
        
        # 2. 添加策略指标（如果策略支持）
        if self.strategy and hasattr(self.strategy, 'get_indicator_values'):
            try:
                indicator_data = self.strategy.get_indicator_values(self.data)
                
                # 根据策略类型添加不同的指标
                if 'MA' in self.strategy.name:
                    # 均线策略
                    for col in indicator_data.columns:
                        if col.startswith('MA'):
                            fig.add_trace(go.Scatter(
                                x=indicator_data.index,
                                y=indicator_data[col],
                                name=col,
                                line=dict(width=1.5),
                            ), row=1, col=1)
            except Exception as e:
                st.warning(f"添加策略指标失败: {e}")
        
        # 3. 标注买入信号
        buy_trades = [t for t in self.result['trades'] if t['类型'] == 'BUY']
        if buy_trades:
            buy_dates = [t['日期'] for t in buy_trades]
            buy_prices = [t['价格'] for t in buy_trades]
            
            fig.add_trace(go.Scatter(
                x=buy_dates,
                y=buy_prices,
                mode='markers',
                name='买入',
                marker=dict(
                    symbol='triangle-up',
                    size=12,
                    color='#EF5350',
                    line=dict(color='white', width=1)
                ),
                hovertemplate='买入<br>日期: %{x}<br>价格: ¥%{y:.2f}<extra></extra>'
            ), row=1, col=1)
        
        # 4. 标注卖出信号
        sell_trades = [t for t in self.result['trades'] if t['类型'] in ['SELL', 'FORCE_CLOSE']]
        if sell_trades:
            sell_dates = [t['日期'] for t in sell_trades]
            sell_prices = [t['价格'] for t in sell_trades]
            
            fig.add_trace(go.Scatter(
                x=sell_dates,
                y=sell_prices,
                mode='markers',
                name='卖出',
                marker=dict(
                    symbol='triangle-down',
                    size=12,
                    color='#26A69A',
                    line=dict(color='white', width=1)
                ),
                hovertemplate='卖出<br>日期: %{x}<br>价格: ¥%{y:.2f}<extra></extra>'
            ), row=1, col=1)
        
        # 5. 持仓数量
        positions_df = self.result['positions']
        fig.add_trace(go.Scatter(
            x=positions_df.index,
            y=positions_df['持仓数量'],
            name='持仓',
            fill='tozeroy',
            fillcolor='rgba(46, 134, 222, 0.3)',
            line=dict(color='#2E86DE', width=2),
        ), row=2, col=1)
        
        # 更新布局
        fig.update_layout(
            height=800,
            showlegend=True,
            hovermode='x unified',
            template='plotly_white',
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        
        # 隐藏rangeslider
        fig.update_xaxes(rangeslider_visible=False, row=1, col=1)
        
        # 设置Y轴标签
        fig.update_yaxes(title_text="价格（元）", row=1, col=1)
        fig.update_yaxes(title_text="数量（股）", row=2, col=1)
        
        return fig
    
    def display_performance_metrics(self, performance: Dict):
        """
        在Streamlit中展示绩效指标
        
        参数:
            performance: 绩效指标字典
        """
        st.subheader("📊 绩效指标")
        
        # 第一行：核心指标
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric(
                "总收益率",
                f"{performance['总收益率']:.2f}%",
                delta=f"{performance['总收益率']:.2f}%" if performance['总收益率'] > 0 else None
            )
        
        with col2:
            st.metric(
                "年化收益率",
                f"{performance['年化收益率']:.2f}%"
            )
        
        with col3:
            st.metric(
                "最大回撤",
                f"{performance['最大回撤']:.2f}%",
                delta=None,
                delta_color="inverse"
            )
        
        with col4:
            st.metric(
                "夏普比率",
                f"{performance['夏普比率']:.2f}"
            )
        
        # 第二行：风险调整收益指标
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric("索提诺比率", f"{performance['索提诺比率']:.2f}")
        
        with col2:
            st.metric("卡玛比率", f"{performance['卡玛比率']:.2f}")
        
        with col3:
            st.metric("波动率", f"{performance['波动率']:.2f}%")
        
        with col4:
            st.metric("下行波动率", f"{performance['下行波动率']:.2f}%")
        
        # 第三行：交易统计
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric("交易次数", f"{performance['交易次数']}")
        
        with col2:
            st.metric("胜率", f"{performance['胜率']:.2f}%")
        
        with col3:
            st.metric("盈亏比", f"{performance['盈亏比']:.2f}")
        
        with col4:
            profit_pct = performance['盈利次数'] / performance['交易次数'] * 100 if performance['交易次数'] > 0 else 0
            st.metric("盈利次数", f"{performance['盈利次数']} ({profit_pct:.1f}%)")
        
        # 详细指标表格
        with st.expander("查看详细指标"):
            col1, col2 = st.columns(2)
            
            with col1:
                st.markdown("**收益指标**")
                st.write(f"- 初始资金: ¥{performance['初始资金']:,.2f}")
                st.write(f"- 最终资金: ¥{performance['最终资金']:,.2f}")
                st.write(f"- 最终现金: ¥{performance['最终现金']:,.2f}")
                st.write(f"- 总收益: ¥{performance['最终资金'] - performance['初始资金']:,.2f}")
                st.write(f"- 回测天数: {performance['回测天数']} 天")
                st.write(f"- 回测年数: {performance['回测年数']:.2f} 年")
            
            with col2:
                st.markdown("**交易统计**")
                st.write(f"- 交易次数: {performance['交易次数']}")
                st.write(f"- 盈利次数: {performance['盈利次数']}")
                st.write(f"- 亏损次数: {performance['亏损次数']}")
                st.write(f"- 平均盈利: ¥{performance['平均盈利']:,.2f}")
                st.write(f"- 平均亏损: ¥{performance['平均亏损']:,.2f}")
                st.write(f"- 最大单笔盈利: ¥{performance['最大单笔盈利']:,.2f}")
                st.write(f"- 最大单笔亏损: ¥{performance['最大单笔亏损']:,.2f}")
    
    def display_trades_table(self, trades: List[Dict]):
        """
        在Streamlit中展示交易记录表格
        
        参数:
            trades: 交易记录列表
        """
        if not trades:
            st.info("暂无交易记录")
            return
        
        st.subheader("📋 交易记录")
        
        # 转换为DataFrame
        trades_df = pd.DataFrame(trades)
        
        # 格式化显示
        if '盈亏' in trades_df.columns:
            trades_df['盈亏'] = trades_df['盈亏'].apply(lambda x: f"¥{x:,.2f}" if pd.notna(x) else "-")
        if '盈亏比例' in trades_df.columns:
            trades_df['盈亏比例'] = trades_df['盈亏比例'].apply(lambda x: f"{x:.2f}%" if pd.notna(x) else "-")
        
        trades_df['价格'] = trades_df['价格'].apply(lambda x: f"¥{x:.2f}")
        trades_df['金额'] = trades_df['金额'].apply(lambda x: f"¥{x:,.2f}")
        trades_df['手续费'] = trades_df['手续费'].apply(lambda x: f"¥{x:.2f}")
        
        # 显示表格
        st.dataframe(
            trades_df,
            use_container_width=True,
            height=400
        )
        
        # 统计摘要
        sell_trades = [t for t in trades if t['类型'] in ['SELL', 'FORCE_CLOSE']]
        if sell_trades:
            col1, col2, col3 = st.columns(3)
            
            total_profit = sum([t['盈亏'] for t in sell_trades if '盈亏' in t])
            win_trades = [t for t in sell_trades if '盈亏' in t and t['盈亏'] > 0]
            loss_trades = [t for t in sell_trades if '盈亏' in t and t['盈亏'] <= 0]
            
            with col1:
                st.metric("总盈亏", f"¥{total_profit:,.2f}")
            
            with col2:
                st.metric("盈利次数/亏损次数", f"{len(win_trades)} / {len(loss_trades)}")
            
            with col3:
                win_rate = len(win_trades) / len(sell_trades) * 100
                st.metric("胜率", f"{win_rate:.2f}%")

